import os

import dotenv
from typing import Annotated

from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import InjectedState, create_react_agent
from langgraph.prebuilt.chat_agent_executor import AgentState

dotenv.load_dotenv()


class CustomState(AgentState):
    user_id: str


@tool(return_direct=True)
def get_user_info(
        state: Annotated[CustomState, InjectedState],
) -> str:
    """查询用户信息"""
    user_id = state["user_id"]
    return "user_123用户的姓名: GuguLH." if user_id == "user_123" else "未知用户"


# 1 构建LLM
llm = ChatOpenAI(
    model="deepseek-chat",
    base_url=os.getenv("DS_BASE"),
    api_key=os.getenv("DS_API_KEY"),
    temperature=0
).bind_tools([get_user_info])

# 2 构建agent
agent = create_react_agent(
    model=llm,
    tools=[get_user_info],
    state_schema=CustomState,
)

# 3 调用agent
print(agent.invoke({
    "messages": "查询用户信息",
    "user_id": "user_123",
}))
